import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger

from transformers import BertTokenizer

from src.models.vec_model.bert_sim_model import BertSimModel


class VectorizeModel:
    def __init__(self, ptm_model_path, device="cpu") -> None:
        self.tokenizer = BertTokenizer.from_pretrained(ptm_model_path)
        self.model = BertSimModel(pretrained_bert_path=ptm_model_path, pooling="cls")
        self.model.eval()

        self.DEVICE = device
        logger.info(device)
        self.model.to(self.DEVICE)

        # 计算两个向量间的欧氏距离：平方和再开方
        self.pdist = nn.PairwiseDistance(2)

    def predict_vec(self, query):
        q_id = self.tokenizer(query, max_length=200, truncation=True, padding="max_length", return_tensors='pt')
        with torch.no_grad():
            q_id_input_ids = q_id["input_ids"].squeeze(1).to(self.DEVICE)
            q_id_attention_mask = q_id["attention_mask"].squeeze(1).to(self.DEVICE)
            q_id_token_type_ids = q_id["token_type_ids"].squeeze(1).to(self.DEVICE)
            q_id_pred = self.model(q_id_input_ids, q_id_attention_mask, q_id_token_type_ids)

        return q_id_pred

    def predict_vec_request(self, query):
        q_id_pred = self.predict_vec(query)
        return q_id_pred.cpu().numpy().tolist()

    def cos_sim(self, q1, q2):
        q1_v = self.predict_vec(q1)
        q2_v = self.predict_vec(q2)
        sim = F.cosine_similarity(q1_v[0], q2_v[0], dim=-1)
        if sim.device.type == 'cuda':
            return sim.cpu().numpy().tolist()
        return sim.numpy().tolist()

    def euclidean_sim(self, q1, q2):
        q1_v = self.predict_vec(q1)
        q2_v = self.predict_vec(q2)
        sim = self.pdist(q1_v, q2_v)
        sim = 1 / (1 + sim)
        if sim.device.type == 'cuda':
            return sim.cpu().numpy()
        return sim.numpy()



if __name__ == "__main__":
    import time, random
    from tqdm import tqdm
    simcse_path = '/home/guweizheng/models/bge-base-en-v1.5'

    vec_model = VectorizeModel(ptm_model_path=simcse_path, device='cuda:0')

    # Simcce
    # print(vec_model.cos_sim("你好啊", "你好"))
    # print(vec_model.cos_sim("你好啊", "你不好"))
    # print(vec_model.cos_sim("你很好", "你不好"))
    # 0.934987485408783
    # 0.8401952385902405
    # 0.9150660037994385

    # print(vec_model.euclidean_sim("你好啊", "你好"))
    # print(vec_model.euclidean_sim("你好啊", "你不好"))
    # print(vec_model.euclidean_sim("你很好", "你不好"))
    # [0.12070419]
    # [0.08026604]
    # [0.10669557]

    # BGE
    # print(vec_model.cos_sim("你好啊", "你好"))
    # print(vec_model.cos_sim("你好啊", "你不好"))
    # print(vec_model.cos_sim("你很好", "你不好"))
    # 0.8723435401916504
    # 0.7223557233810425
    # 0.7223557233810425

    print(vec_model.euclidean_sim("你好啊", "你好"))
    print(vec_model.euclidean_sim("你好啊", "你不好"))
    print(vec_model.euclidean_sim("你很好", "你不好"))
    # [0.11970119]
    # [0.08388899]
    # [0.08388899]


    # print(vec_model.predict_vec("什么人不能吃花生").shape)
    # vec_model.predict_vec("你好啊")
    # vec_model.predict_vec("你好啊")
    # print(vec_model.predict_vec("你好啊"))
    # print(type(vec_model.predict_vec("你好啊")))

    # tmp_queries = ["你好啊", "今天天气怎么样", "我要暴富"]
    # for i in tqdm(range(200)):
    #     vec_model.predict_vec(random.choice(tmp_queries))
    # for i in tqdm(range(200)):
    #     vec_model.predict_vec(random.choice(tmp_queries))
    # for i in range(100):
    #     print(i)
    #     time.sleep(1)
